// Copyright (C) 2017  Miroslav Lichvar
// Copyright (C) 2021  Varphone Wong
//
// This program is free software; you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation; either version 2 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

use byteorder::{BigEndian, ByteOrder};
use chrono::NaiveDateTime;
use log::{debug, info, warn};
use rand::random;
use socket2::{Domain, Socket, Type};
use std::io;
use std::io::{Error, ErrorKind};
use std::net::{SocketAddr, ToSocketAddrs, UdpSocket};
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::{Arc, Mutex};
use std::thread::{self, JoinHandle};
use std::time::{Duration, SystemTime};

const NTP_EPOCH_SECS: u64 = 2_208_988_800;
const NANO_UNIT: u64 = 1_000_000_000;
const READ_TMO_MS: u64 = 40;

/// 一个描述 NTP 时戳的类型。
#[derive(Debug, Copy, Clone)]
pub struct NtpTimestamp {
    ts: u64,
}

impl NtpTimestamp {
    /// 返回当前时间的 NTP 时戳。
    pub fn now() -> NtpTimestamp {
        let now = SystemTime::now();
        let dur = now.duration_since(std::time::UNIX_EPOCH).unwrap();
        let secs = dur.as_secs() + NTP_EPOCH_SECS;
        let nanos = dur.subsec_nanos();

        NtpTimestamp {
            ts: (secs << 32) + (nanos as u64 * u32::MAX as u64 / NANO_UNIT),
        }
    }

    /// 返回 “零点” 的 NTP 时戳。
    pub fn zero() -> NtpTimestamp {
        NtpTimestamp { ts: 0 }
    }

    /// 返回随机 NTP 时戳。
    pub fn random() -> NtpTimestamp {
        NtpTimestamp { ts: random() }
    }

    /// 返回与指定时戳的差值。
    pub fn diff_to_sec(&self, ts: &NtpTimestamp) -> f64 {
        (self.ts - ts.ts) as i64 as f64 / 4294967296.0
    }

    /// 返回秒数部分。
    pub fn subsec_secs(&self) -> u64 {
        self.ts >> 32
    }

    /// 返回自 1970/1/1 起的秒数部分。
    pub fn subsec_secs_epoch(&self) -> u64 {
        self.subsec_secs().saturating_sub(NTP_EPOCH_SECS)
    }

    /// 返回纳秒数部分。
    pub fn subsec_nanos(&self) -> u64 {
        ((self.ts & 0x0000_0000_ffff_ffff) * u32::MAX as u64) / NANO_UNIT
    }

    /// 从字节流读入一个时戳。
    pub fn read(buf: &[u8]) -> NtpTimestamp {
        NtpTimestamp {
            ts: BigEndian::read_u64(buf),
        }
    }

    /// 将时戳写到字节流中。
    pub fn write(&self, buf: &mut [u8]) {
        BigEndian::write_u64(buf, self.ts);
    }
}

impl From<&AtomicU64> for NtpTimestamp {
    fn from(val: &AtomicU64) -> Self {
        Self {
            ts: val.load(Ordering::SeqCst),
        }
    }
}

impl From<NtpTimestamp> for u64 {
    fn from(val: NtpTimestamp) -> Self {
        val.ts
    }
}

impl TryFrom<NaiveDateTime> for NtpTimestamp {
    type Error = &'static str;
    fn try_from(val: NaiveDateTime) -> Result<Self, Self::Error> {
        let secs = val.timestamp() as u64 + NTP_EPOCH_SECS;
        if secs <= u32::MAX as u64 {
            let nanos = (val.timestamp_subsec_nanos() as u64 * u32::MAX as u64) / NANO_UNIT;
            Ok(Self {
                ts: (secs << 32) + nanos,
            })
        } else {
            Err("NTP timestamp cannot excced 2036-02-07T06:28:15")
        }
    }
}

impl From<SystemTime> for NtpTimestamp {
    fn from(val: SystemTime) -> Self {
        let dur = val
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap_or_else(|_| Duration::from_secs(0));
        let secs = dur.as_secs() + NTP_EPOCH_SECS;
        let nanos = dur.subsec_nanos() as u64;

        Self {
            ts: (secs << 32) + (nanos * u32::MAX as u64) / NANO_UNIT,
        }
    }
}

impl From<NtpTimestamp> for NaiveDateTime {
    fn from(val: NtpTimestamp) -> Self {
        NaiveDateTime::from_timestamp(val.subsec_secs_epoch() as i64, val.subsec_nanos() as u32)
    }
}

impl From<NtpTimestamp> for SystemTime {
    fn from(val: NtpTimestamp) -> Self {
        use std::time::UNIX_EPOCH;

        let secs = (val.ts >> 32) - NTP_EPOCH_SECS;
        let nanos = ((val.ts & 0x0000_0000_ffff_ffff) as f64 / 4.294967296) as u64;
        UNIX_EPOCH
            .checked_add(Duration::from_secs(secs))
            .and_then(|x| x.checked_add(Duration::from_nanos(nanos)))
            .unwrap_or(UNIX_EPOCH)
    }
}

impl PartialEq for NtpTimestamp {
    fn eq(&self, other: &NtpTimestamp) -> bool {
        self.ts == other.ts
    }
}

#[derive(Debug, Copy, Clone)]
pub struct NtpFracValue {
    val: u32,
}

impl NtpFracValue {
    pub fn read(buf: &[u8]) -> NtpFracValue {
        NtpFracValue {
            val: BigEndian::read_u32(buf),
        }
    }

    pub fn write(&self, buf: &mut [u8]) {
        BigEndian::write_u32(buf, self.val);
    }

    pub fn zero() -> NtpFracValue {
        NtpFracValue { val: 0 }
    }

    #[allow(dead_code)]
    pub fn increment(&mut self) {
        self.val += 1;
    }
}

/// 一个描述 NTP 数据包的类型。
#[derive(Debug)]
pub struct NtpPacket {
    remote_addr: SocketAddr,
    local_ts: NtpTimestamp,

    leap: u8,
    version: u8,
    mode: u8,
    stratum: u8,
    poll: i8,
    precision: i8,
    delay: NtpFracValue,
    dispersion: NtpFracValue,
    ref_id: u32,
    ref_ts: NtpTimestamp,
    orig_ts: NtpTimestamp,
    rx_ts: NtpTimestamp,
    tx_ts: NtpTimestamp,
}

impl NtpPacket {
    /// 从指定 UDP 套接字中接收一个 NTP 数据包。
    pub fn receive(socket: &UdpSocket) -> io::Result<NtpPacket> {
        let mut buf = [0; 1024];

        let (len, addr) = socket.recv_from(&mut buf)?;

        let local_ts = NtpTimestamp::now();

        if len < 48 {
            return Err(Error::new(ErrorKind::UnexpectedEof, "数据包太短"));
        }

        let leap = buf[0] >> 6;
        let version = (buf[0] >> 3) & 0x7;
        let mode = buf[0] & 0x7;

        if !(1..=4).contains(&version) {
            return Err(Error::new(ErrorKind::Other, "不支持的版本"));
        }

        Ok(NtpPacket {
            remote_addr: addr,
            local_ts,
            leap,
            version,
            mode,
            stratum: buf[1],
            poll: buf[2] as i8,
            precision: buf[3] as i8,
            delay: NtpFracValue::read(&buf[4..8]),
            dispersion: NtpFracValue::read(&buf[8..12]),
            ref_id: BigEndian::read_u32(&buf[12..16]),
            ref_ts: NtpTimestamp::read(&buf[16..24]),
            orig_ts: NtpTimestamp::read(&buf[24..32]),
            rx_ts: NtpTimestamp::read(&buf[32..40]),
            tx_ts: NtpTimestamp::read(&buf[40..48]),
        })
    }

    /// 将当前 NTP 数据包发送到 UDP 套接字中。
    pub fn send(&self, socket: &UdpSocket) -> io::Result<usize> {
        let mut buf = [0; 48];

        buf[0] = self.leap << 6 | self.version << 3 | self.mode;
        buf[1] = self.stratum;
        buf[2] = self.poll as u8;
        buf[3] = self.precision as u8;
        self.delay.write(&mut buf[4..8]);
        self.dispersion.write(&mut buf[8..12]);
        BigEndian::write_u32(&mut buf[12..16], self.ref_id);
        self.ref_ts.write(&mut buf[16..24]);
        self.orig_ts.write(&mut buf[24..32]);
        self.rx_ts.write(&mut buf[32..40]);
        self.tx_ts.write(&mut buf[40..48]);

        socket.send_to(&buf, self.remote_addr)
    }

    /// 当为请求数据包时返回 true。
    pub fn is_request(&self) -> bool {
        self.mode == 1
            || self.mode == 3
            || (self.mode == 0 && self.version == 1 && self.remote_addr.port() != 123)
    }

    /// 为当前请求生成回应数据包。
    pub fn make_response(&self, state: &NtpServerState) -> Option<NtpPacket> {
        if !self.is_request() {
            return None;
        }

        Some(NtpPacket {
            remote_addr: self.remote_addr,
            local_ts: NtpTimestamp::zero(),
            leap: state.leap,
            version: self.version,
            mode: if self.mode == 1 { 2 } else { 4 },
            stratum: state.stratum,
            poll: self.poll,
            precision: state.precision,
            delay: state.delay,
            dispersion: state.dispersion,
            ref_id: state.ref_id,
            ref_ts: state.ref_ts,
            orig_ts: self.tx_ts,
            rx_ts: self.local_ts,
            tx_ts: NtpTimestamp::now(),
        })
    }

    /// 生成一份新的请求数据包。
    pub fn new_request(remote_addr: SocketAddr) -> NtpPacket {
        NtpPacket {
            remote_addr,
            local_ts: NtpTimestamp::now(),
            leap: 0,
            version: 4,
            mode: 3,
            stratum: 0,
            poll: 0,
            precision: 0,
            delay: NtpFracValue::zero(),
            dispersion: NtpFracValue::zero(),
            ref_id: 0,
            ref_ts: NtpTimestamp::zero(),
            orig_ts: NtpTimestamp::zero(),
            rx_ts: NtpTimestamp::zero(),
            tx_ts: NtpTimestamp::random(),
        }
    }

    /// 当为回应数据包时返回 true。
    pub fn is_valid_response(&self, request: &NtpPacket) -> bool {
        self.remote_addr == request.remote_addr
            && self.mode == request.mode + 1
            && self.orig_ts == request.tx_ts
    }

    /// 返回当前数据包的服务器状态。
    pub fn get_server_state(&self) -> NtpServerState {
        NtpServerState {
            leap: self.leap,
            stratum: self.stratum,
            precision: self.precision,
            ref_id: self.ref_id,
            ref_ts: self.ref_ts,
            dispersion: self.dispersion,
            delay: self.delay,
        }
    }
}

/// 一个描述 NTP 服务器状态的类型。
#[derive(Copy, Clone, Debug)]
pub struct NtpServerState {
    leap: u8,
    stratum: u8,
    precision: i8,
    ref_id: u32,
    ref_ts: NtpTimestamp,
    dispersion: NtpFracValue,
    delay: NtpFracValue,
}

impl NtpServerState {
    /// 返回一份空的服务器状态。
    pub fn with_empty() -> Self {
        Self {
            leap: 0,
            stratum: 0,
            precision: 0,
            ref_id: 0,
            ref_ts: NtpTimestamp::zero(),
            dispersion: NtpFracValue::zero(),
            delay: NtpFracValue::zero(),
        }
    }

    /// 使用参考时钟构建一份服务器状态。
    pub fn with_ref_ts(ref_ts: &AtomicU64) -> Self {
        Self {
            leap: 0,
            stratum: 2,
            precision: 0,
            ref_id: 0,
            ref_ts: NtpTimestamp::from(ref_ts),
            dispersion: NtpFracValue { val: 0 },
            delay: NtpFracValue { val: 100 },
        }
    }
}

/// 一个描述 NTP 服务器的类型。
pub struct NtpServer {
    sockets: Vec<Arc<UdpSocket>>,
    exit_loop: Arc<AtomicBool>,
    ref_ts: Arc<AtomicU64>,
    threads: Arc<Mutex<Vec<JoinHandle<()>>>>,
}

impl NtpServer {
    /// 创建一个 NTP 服务器实例。
    pub fn new<A: ToSocketAddrs>(addrs: &[A]) -> Result<NtpServer, io::Error> {
        let tmo = Duration::from_millis(READ_TMO_MS);
        let mut sockets = vec![];

        for s in addrs {
            for a in s.to_socket_addrs()? {
                let socket = Socket::new(Domain::for_address(a), Type::DGRAM, None)?;
                socket.set_read_timeout(Some(tmo))?;
                socket.set_reuse_address(true)?;
                socket.set_reuse_port(true)?;
                socket.bind(&a.into())?;
                sockets.push(Arc::new(socket.into()));
            }
        }

        Ok(NtpServer {
            sockets,
            exit_loop: Arc::new(AtomicBool::new(false)),
            ref_ts: Arc::new(AtomicU64::new(0)),
            threads: Arc::new(Mutex::new(vec![])),
        })
    }

    fn process_requests(
        id: usize,
        socket: Arc<UdpSocket>,
        exit_loop: Arc<AtomicBool>,
        ref_ts: Arc<AtomicU64>,
    ) {
        let mut last_update = NtpTimestamp::now();
        let mut cached_state: NtpServerState;
        cached_state = NtpServerState::with_ref_ts(&ref_ts);

        info!(
            "#{} 开始监听 {:?} 提供网络时钟服务",
            id,
            socket.local_addr().unwrap()
        );

        while !exit_loop.load(Ordering::SeqCst) {
            match NtpPacket::receive(&socket) {
                Ok(request) => {
                    debug!("#{} 收到来自 {} 的请求", id, request.remote_addr);

                    if request.local_ts.diff_to_sec(&last_update).abs() > 0.1 {
                        cached_state = NtpServerState::with_ref_ts(&ref_ts);
                        last_update = request.local_ts;
                    }

                    if let Some(response) = request.make_response(&cached_state) {
                        match response.send(&socket) {
                            Ok(_) => {}
                            Err(e) => {
                                warn!("#{} 回应 {} 时发生异常：{}", id, response.remote_addr, e)
                            }
                        }
                    }
                }
                Err(e) => {
                    if e.kind() == ErrorKind::WouldBlock {
                        //
                    } else {
                        warn!("#{} 接收数据包时发生异常：{}", id, e);
                    }
                }
            }
        }
    }

    /// 等待后他线程退出。
    pub fn join(&self) {
        if let Ok(mut threads) = self.threads.lock() {
            while let Some(thread) = threads.pop() {
                let _ = thread.join();
            }
        }
    }

    /// 运行服务并阻塞当前线程。
    pub fn run(&self) {
        self.spawn();
        self.join();
    }

    /// 关闭服务。
    pub fn shutdown(&self) {
        self.exit_loop.store(true, Ordering::SeqCst);
    }

    /// 置于后台运行。
    pub fn spawn(&self) {
        if let Ok(mut threads) = self.threads.lock() {
            for (id, socket) in self.sockets.iter().enumerate() {
                let socket = Arc::clone(socket);
                let exit_loop = Arc::clone(&self.exit_loop);
                let ref_ts = Arc::clone(&self.ref_ts);
                threads.push(thread::spawn(move || {
                    NtpServer::process_requests(id, socket, exit_loop, ref_ts);
                }));
            }
        }
    }

    /// 更新参考时钟。
    pub fn update_ref_ts(&self, ts: NtpTimestamp) {
        self.ref_ts.store(ts.into(), Ordering::SeqCst);
    }
}

impl Drop for NtpServer {
    fn drop(&mut self) {
        self.shutdown();
        self.join();
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use chrono::NaiveDate;

    #[test]
    fn test_ntp_timestamp() {
        let src_dt = NaiveDate::from_ymd(1970, 1, 1).and_hms(0, 0, 1);
        assert_eq!(src_dt.timestamp_nanos(), 1_000_000_000);
        let ntp_ts = NtpTimestamp::try_from(src_dt).unwrap();
        assert_eq!(ntp_ts.subsec_secs(), NTP_EPOCH_SECS + 1);
        let dst_dt = NaiveDateTime::from(ntp_ts);
        assert_eq!(src_dt, dst_dt);
        for y in 1970..=2036 {
            let src_dt = NaiveDate::from_ymd(y, 1, 1).and_hms(0, 0, 0);
            let ntp_ts = NtpTimestamp::try_from(src_dt).unwrap();
            let dst_dt = NaiveDateTime::from(ntp_ts);
            assert_eq!(src_dt, dst_dt);
        }
        let epoch_dt = NaiveDate::from_ymd(1970, 1, 1).and_hms(0, 0, 0);
        // 1980-01-01T00:00:00 ~ 2036-02-07T06:28:15
        for d in 0..24_143 {
            if let Some(src_dt) = epoch_dt.checked_add_signed(chrono::Duration::days(d)) {
                let ntp_ts = NtpTimestamp::try_from(src_dt).unwrap();
                let dst_dt = NaiveDateTime::from(ntp_ts);
                assert_eq!(src_dt, dst_dt);
            }
        }
    }

    #[test]
    #[should_panic]
    fn test_ntp_timestamp_exceed() {
        let src_dt = NaiveDate::from_ymd(2036, 2, 7).and_hms(6, 28, 16);
        let _ntp_ts = NtpTimestamp::try_from(src_dt).unwrap();
    }
}
